// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <gmock/gmock.h>
#include <gtest/gtest.h>

#include <pollux/expression/function_signature.h>
#include <pollux/expression/register_special_form.h>
#include <pollux/expression/vector_function.h>
#include <pollux/functions/function_registry.h>
#include <pollux/functions/macros.h>
#include <pollux/functions/registerer.h>
#include <pollux/functions/prestosql/aggregates/register_aggregate_functions.h>
#include <pollux/functions/prestosql/registration/registration_functions.h>
#include <pollux/testing/functions/prestosql/function_base_test.h>
#include <pollux/functions/prestosql/types/ip_prefix_registration.h>
#include <pollux/functions/prestosql/types/ip_prefix_type.h>
#include <tests/functions/registry_test_util.h>
#include <pollux/type/type.h>

namespace kumo::pollux {

namespace {

POLLUX_DECLARE_VECTOR_FUNCTION(
    udf_vector_func_one,
    VectorFuncOne::signatures(),
    std::make_unique<VectorFuncOne>());

POLLUX_DECLARE_VECTOR_FUNCTION(
    udf_vector_func_two,
    VectorFuncTwo::signatures(),
    std::make_unique<VectorFuncTwo>());

POLLUX_DECLARE_VECTOR_FUNCTION(
    udf_vector_func_three,
    VectorFuncThree::signatures(),
    std::make_unique<VectorFuncThree>());

POLLUX_DECLARE_VECTOR_FUNCTION_WITH_METADATA(
    udf_vector_func_four,
    VectorFuncFour::signatures(),
    exec::VectorFunctionMetadataBuilder().deterministic(false).build(),
    std::make_unique<VectorFuncFour>());

inline void registerTestFunctions() {
  // If no alias is specified, ensure it will fallback to the struct name.
  register_function<FuncOne, Varchar, Varchar>({"func_one", "Func_One_Alias"});

  // func_two has two signatures.
  register_function<FuncTwo, int64_t, int64_t, int32_t>({"func_two"});
  register_function<FuncTwo, int64_t, int64_t, int16_t>({"func_two"});

  // func_three has two aliases.
  register_function<FuncThree, Array<int64_t>, Array<int64_t>>(
      {"func_three_alias1", "func_three_alias2"});

  // We swap func_four and func_five while registering.
  register_function<FuncFour, Varchar, Varchar>({"func_five"});
  register_function<FuncFive, int64_t, int64_t>({"func_four"});

  register_function<VariadicFunc, Varchar, Variadic<Varchar>>(
      {"variadic_func", "Variadic_Func_Alias"});

  POLLUX_REGISTER_VECTOR_FUNCTION(udf_vector_func_one, "vector_func_one");
  POLLUX_REGISTER_VECTOR_FUNCTION(udf_vector_func_one, "Vector_Func_One_Alias");
  POLLUX_REGISTER_VECTOR_FUNCTION(udf_vector_func_two, "vector_func_two");
  POLLUX_REGISTER_VECTOR_FUNCTION(udf_vector_func_three, "vector_func_three");
  POLLUX_REGISTER_VECTOR_FUNCTION(udf_vector_func_four, "vector_func_four");
}

inline void registerTestVectorFunctionOne(const std::string& functionName) {
  POLLUX_REGISTER_VECTOR_FUNCTION(udf_vector_func_one, functionName);
}
} // namespace

class FunctionRegistryTest : public testing::Test {
 public:
  FunctionRegistryTest() {
    registerTestFunctions();
    exec::registerFunctionCallToSpecialForms();
  }

  void testResolveVectorFunction(
      const std::string& functionName,
      const std::vector<TypePtr>& types,
      const TypePtr& expected) {
    checkEqual(pollux::resolve_function(functionName, types), expected);
    checkEqual(pollux::resolveVectorFunction(functionName, types), expected);
  }

  void checkEqual(const TypePtr& actual, const TypePtr& expected) {
    if (expected) {
      EXPECT_EQ(*actual, *expected);
    } else {
      EXPECT_EQ(actual, nullptr);
    }
  }
};

TEST_F(FunctionRegistryTest, removeFunction) {
  const std::string functionName = "func_to_remove";
  auto checkFunctionExists = [&](const std::string& name,
                                 bool vectorFuncSignatures,
                                 bool simpleFuncSignatures) {
    EXPECT_EQ(
        get_function_signatures(name).size(),
        vectorFuncSignatures + simpleFuncSignatures);
    EXPECT_EQ(getVectorFunctionSignatures().count(name), vectorFuncSignatures);
    EXPECT_EQ(
        exec::simple_functions().get_function_signatures(name).size(),
        simpleFuncSignatures);
  };

  checkFunctionExists(functionName, 0, 0);

  // Only vector function registered
  registerTestVectorFunctionOne(functionName);
  checkFunctionExists(functionName, 1, 0);
  removeFunction(functionName);
  checkFunctionExists(functionName, 0, 0);

  // Only simple function registered
  register_function<FuncOne, Varchar, Varchar>(
      std::vector<std::string>{functionName});
  checkFunctionExists(functionName, 0, 1);
  removeFunction(functionName);
  checkFunctionExists(functionName, 0, 0);

  // Both vector and simple function registered
  registerTestVectorFunctionOne(functionName);
  register_function<FuncOne, Varchar, Varchar>(
      std::vector<std::string>{functionName});
  checkFunctionExists(functionName, 1, 1);
  removeFunction(functionName);
  checkFunctionExists(functionName, 0, 0);
}

TEST_F(FunctionRegistryTest, getFunctionSignaturesByName) {
  {
    auto signatures = get_function_signatures("func_one");
    ASSERT_EQ(signatures.size(), 1);
    ASSERT_EQ(
        signatures.at(0)->toString(),
        exec::FunctionSignatureBuilder()
            .returnType("varchar")
            .argumentType("varchar")
            .build()
            ->toString());
  }

  {
    auto signatures = get_function_signatures("vector_func_one");
    ASSERT_EQ(signatures.size(), 1);
    ASSERT_EQ(
        signatures.at(0)->toString(),
        exec::FunctionSignatureBuilder()
            .returnType("bigint")
            .argumentType("varchar")
            .build()
            ->toString());
  }

  ASSERT_TRUE(get_function_signatures("non-existent-function").empty());
}

TEST_F(FunctionRegistryTest, get_function_signatures) {
  auto functionSignatures = get_function_signatures();
  ASSERT_EQ(functionSignatures.size(), 14);

  ASSERT_EQ(functionSignatures.count("func_one"), 1);
  ASSERT_EQ(functionSignatures.count("func_two"), 1);
  ASSERT_EQ(functionSignatures.count("func_three"), 0);
  ASSERT_EQ(functionSignatures.count("func_three_alias1"), 1);
  ASSERT_EQ(functionSignatures.count("func_three_alias2"), 1);
  ASSERT_EQ(functionSignatures.count("func_four"), 1);
  ASSERT_EQ(functionSignatures.count("func_five"), 1);
  ASSERT_EQ(functionSignatures.count("variadic_func"), 1);
  ASSERT_EQ(functionSignatures.count("vector_func_one"), 1);
  ASSERT_EQ(functionSignatures.count("vector_func_two"), 1);
  ASSERT_EQ(functionSignatures.count("vector_func_three"), 1);
  ASSERT_EQ(functionSignatures.count("vector_func_four"), 1);

  ASSERT_EQ(functionSignatures["func_one"].size(), 1);
  ASSERT_EQ(functionSignatures["func_two"].size(), 2);
  ASSERT_EQ(functionSignatures["func_three"].size(), 0);
  ASSERT_EQ(functionSignatures["func_three_alias1"].size(), 1);
  ASSERT_EQ(functionSignatures["func_three_alias2"].size(), 1);
  ASSERT_EQ(functionSignatures["vector_func_one"].size(), 1);
  ASSERT_EQ(functionSignatures["vector_func_two"].size(), 1);
  ASSERT_EQ(functionSignatures["vector_func_three"].size(), 1);
  ASSERT_EQ(functionSignatures["vector_func_four"].size(), 1);

  ASSERT_EQ(
      functionSignatures["func_one"].at(0)->toString(),
      exec::FunctionSignatureBuilder()
          .returnType("varchar")
          .argumentType("varchar")
          .build()
          ->toString());

  std::vector<std::string> funcTwoSignatures;
  std::transform(
      functionSignatures["func_two"].begin(),
      functionSignatures["func_two"].end(),
      std::back_inserter(funcTwoSignatures),
      [](auto& signature) { return signature->toString(); });
  ASSERT_THAT(
      funcTwoSignatures,
      testing::UnorderedElementsAre(
          exec::FunctionSignatureBuilder()
              .returnType("bigint")
              .argumentType("bigint")
              .argumentType("integer")
              .build()
              ->toString(),
          exec::FunctionSignatureBuilder()
              .returnType("bigint")
              .argumentType("bigint")
              .argumentType("smallint")
              .build()
              ->toString()));

  ASSERT_EQ(
      functionSignatures["func_three_alias1"].at(0)->toString(),
      exec::FunctionSignatureBuilder()
          .returnType("array(bigint)")
          .argumentType("array(bigint)")
          .build()
          ->toString());

  ASSERT_EQ(
      functionSignatures["func_three_alias2"].at(0)->toString(),
      exec::FunctionSignatureBuilder()
          .returnType("array(bigint)")
          .argumentType("array(bigint)")
          .build()
          ->toString());

  ASSERT_EQ(
      functionSignatures["func_four"].at(0)->toString(),
      exec::FunctionSignatureBuilder()
          .returnType("bigint")
          .argumentType("bigint")
          .build()
          ->toString());

  ASSERT_EQ(
      functionSignatures["func_five"].at(0)->toString(),
      exec::FunctionSignatureBuilder()
          .returnType("varchar")
          .argumentType("varchar")
          .build()
          ->toString());

  ASSERT_EQ(
      functionSignatures["variadic_func"].at(0)->toString(),
      exec::FunctionSignatureBuilder()
          .returnType("varchar")
          .variableArity("varchar")
          .build()
          ->toString());

  ASSERT_EQ(
      functionSignatures["vector_func_one"].at(0)->toString(),
      exec::FunctionSignatureBuilder()
          .returnType("bigint")
          .argumentType("varchar")
          .build()
          ->toString());

  ASSERT_EQ(
      functionSignatures["vector_func_two"].at(0)->toString(),
      exec::FunctionSignatureBuilder()
          .returnType("array(bigint)")
          .argumentType("array(varchar)")
          .build()
          ->toString());

  ASSERT_EQ(
      functionSignatures["vector_func_three"].at(0)->toString(),
      exec::FunctionSignatureBuilder()
          .returnType("opaque")
          .argumentType("any")
          .build()
          ->toString());

  ASSERT_EQ(
      functionSignatures["vector_func_four"].at(0)->toString(),
      exec::FunctionSignatureBuilder()
          .knownTypeVariable("K")
          .typeVariable("V")
          .returnType("array(K)")
          .argumentType("map(K,V)")
          .build()
          ->toString());
}

TEST_F(FunctionRegistryTest, getVectorFunctionSignatures) {
  auto functionSignatures = getVectorFunctionSignatures();
  ASSERT_EQ(functionSignatures.size(), 5);

  std::set<std::string> functionNames;
  std::transform(
      functionSignatures.begin(),
      functionSignatures.end(),
      std::inserter(functionNames, functionNames.end()),
      [](auto& signature) { return signature.first; });

  ASSERT_THAT(
      functionNames,
      ::testing::UnorderedElementsAre(
          "vector_func_one",
          "vector_func_one_alias",
          "vector_func_two",
          "vector_func_three",
          "vector_func_four"));
}

TEST_F(FunctionRegistryTest, hasSimpleFunctionSignature) {
  auto result = resolve_function("func_one", {VARCHAR()});
  ASSERT_EQ(*result, *VARCHAR());
}

TEST_F(FunctionRegistryTest, hasSimpleFunctionSignatureWrongArgType) {
  auto result = resolve_function("func_one", {INTEGER()});
  ASSERT_EQ(result, nullptr);
}

TEST_F(FunctionRegistryTest, hasSimpleFunctionSignatureWrongFunctionName) {
  auto result = resolve_function("method_one", {VARCHAR()});
  ASSERT_EQ(result, nullptr);
}

TEST_F(FunctionRegistryTest, hasVariadicFunctionSignature) {
  auto result = resolve_function("variadic_func", {});
  ASSERT_EQ(*result, *VARCHAR());

  result = resolve_function("variadic_func", {VARCHAR()});
  ASSERT_EQ(*result, *VARCHAR());

  result = resolve_function("variadic_func", {VARCHAR(), VARCHAR()});
  ASSERT_EQ(*result, *VARCHAR());

  result = resolve_function("variadic_func", {INTEGER()});
  ASSERT_EQ(result, nullptr);
}

TEST_F(FunctionRegistryTest, hasVectorFunctionSignature) {
  testResolveVectorFunction("vector_func_one", {VARCHAR()}, BIGINT());
}

TEST_F(FunctionRegistryTest, hasVectorFunctionSignature2) {
  testResolveVectorFunction(
      "vector_func_two", {ARRAY(VARCHAR())}, ARRAY(BIGINT()));
}

TEST_F(FunctionRegistryTest, hasVectorFunctionSignature3) {
  testResolveVectorFunction("vector_func_three", {REAL()}, OPAQUE<void>());
}

TEST_F(FunctionRegistryTest, hasVectorFunctionSignature4) {
  testResolveVectorFunction(
      "vector_func_four", {MAP(BIGINT(), VARCHAR())}, ARRAY(BIGINT()));
}

TEST_F(FunctionRegistryTest, hasVectorFunctionSignatureWrongArgType) {
  testResolveVectorFunction("vector_func_one", {INTEGER()}, nullptr);
}

TEST_F(FunctionRegistryTest, hasVectorFunctionSignatureWrongFunctionName) {
  testResolveVectorFunction("vector_method_one", {VARCHAR()}, nullptr);
}

TEST_F(FunctionRegistryTest, registerFunctionTwice) {
  // For better or worse, there are code paths that depend on the ability to
  // register the same functions repeatedly and have those repeated calls
  // ignored.
  register_function<FuncOne, Varchar, Varchar>({"func_one"});
  register_function<FuncOne, Varchar, Varchar>({"func_one"});

  auto& functions = exec::simple_functions();
  auto signatures = functions.get_function_signatures("func_one");
  // The function should only be registered once, despite the multiple calls to
  // register_function.
  ASSERT_EQ(signatures.size(), 1);
}

TEST_F(FunctionRegistryTest, functionNameInMixedCase) {
  auto result = resolve_function("funC_onE", {VARCHAR()});
  ASSERT_EQ(*result, *VARCHAR());
  result = resolve_function("funC_onE_aliaS", {VARCHAR()});
  ASSERT_EQ(*result, *VARCHAR());

  testResolveVectorFunction("vectoR_funC_onE_aliaS", {VARCHAR()}, BIGINT());
  testResolveVectorFunction("vectoR_funC_onE", {VARCHAR()}, BIGINT());

  result = resolve_function("variadiC_funC_aliaS", {VARCHAR(), VARCHAR()});
  ASSERT_EQ(*result, *VARCHAR());
  result = resolve_function("variadiC_funC", {});
  ASSERT_EQ(*result, *VARCHAR());
}

TEST_F(FunctionRegistryTest, isDeterministic) {
  functions::prestosql::registerAllScalarFunctions();
  ASSERT_TRUE(isDeterministic("plus").value());
  ASSERT_TRUE(isDeterministic("in").value());

  ASSERT_FALSE(isDeterministic("rand").value());
  ASSERT_FALSE(isDeterministic("uuid").value());
  ASSERT_FALSE(isDeterministic("shuffle").value());

  // Not found functions.
  ASSERT_FALSE(isDeterministic("cast").has_value());
  ASSERT_FALSE(isDeterministic("not_found_function").has_value());
}

TEST_F(FunctionRegistryTest, companionFunction) {
  functions::prestosql::registerAllScalarFunctions();
  aggregate::prestosql::registerAllAggregateFunctions();
  const auto functions = {"array_frequency", "bitwise_left_shift", "ceil"};
  // Aggregate companion functions with suffix '_extract' are registered as
  // vector functions.
  const auto companionFunctions = {
      "array_agg_extract", "arbitrary_extract", "bitwise_and_agg_extract"};

  for (const auto& function : functions) {
    ASSERT_FALSE(exec::simple_functions()
                     .get_function_signatures_and_metadata(function)
                     .front()
                     .first.companionFunction);
  }
  for (const auto& function : companionFunctions) {
    ASSERT_TRUE(exec::getVectorFunctionMetadata(function)->companionFunction);
  }
}

template <typename T>
struct TestFunction {
  POLLUX_DEFINE_FUNCTION_TYPES(T);

  void call(
      out_type<Varchar>& out,
      const arg_type<Varchar>&,
      const arg_type<Varchar>&) {
    out.copy_from("1"_sv);
  }

  void call(int32_t& out, const arg_type<Variadic<Varchar>>&) {
    out = 2;
  }

  void
  call(float& out, const arg_type<Generic<T1>>&, const arg_type<Generic<T1>>&) {
    out = 3;
  }

  void call(int64_t& out, const arg_type<Variadic<Any>>&) {
    out = 4;
  }

  void
  call(double& out, const arg_type<Varchar>&, const arg_type<Variadic<Any>>&) {
    out = 5;
  }
};

TEST_F(FunctionRegistryTest, resolveFunctionsBasedOnPriority) {
  std::string func = "func_with_priority";

  register_function<TestFunction, double, Varchar, Variadic<Any>>({func});
  register_function<TestFunction, Varchar, Varchar, Varchar>({func});
  register_function<TestFunction, int64_t, Variadic<Any>>({func});
  register_function<TestFunction, int32_t, Variadic<Varchar>>({func});
  register_function<TestFunction, float, Generic<T1>, Generic<T1>>({func});

  auto result1 = resolve_function(func, {VARCHAR(), VARCHAR()});
  ASSERT_EQ(*result1, *VARCHAR());

  auto result2 = resolve_function(func, {VARCHAR(), VARCHAR(), VARCHAR()});
  ASSERT_EQ(*result2, *INTEGER());

  auto result3 = resolve_function(func, {VARCHAR(), INTEGER()});
  ASSERT_EQ(*result3, *DOUBLE());

  auto result4 = resolve_function(func, {INTEGER(), VARCHAR()});
  ASSERT_EQ(*result4, *BIGINT());

  auto result5 = resolve_function(func, {INTEGER(), INTEGER()});
  ASSERT_EQ(*result5, *REAL());
}

TEST_F(FunctionRegistryTest, resolveSpecialForms) {
  auto andResult =
      resolveFunctionOrCallableSpecialForm("and", {BOOLEAN(), BOOLEAN()});
  ASSERT_EQ(*andResult, *BOOLEAN());

  auto coalesceResult =
      resolveFunctionOrCallableSpecialForm("coalesce", {VARCHAR(), VARCHAR()});
  ASSERT_EQ(*coalesceResult, *VARCHAR());

  auto ifResult = resolveFunctionOrCallableSpecialForm(
      "if", {BOOLEAN(), INTEGER(), INTEGER()});
  ASSERT_EQ(*ifResult, *INTEGER());

  auto orResult =
      resolveFunctionOrCallableSpecialForm("or", {BOOLEAN(), BOOLEAN()});
  ASSERT_EQ(*orResult, *BOOLEAN());

  auto switchResult = resolveFunctionOrCallableSpecialForm(
      "switch", {BOOLEAN(), DOUBLE(), BOOLEAN(), DOUBLE(), DOUBLE()});
  ASSERT_EQ(*switchResult, *DOUBLE());

  auto tryResult = resolveFunctionOrCallableSpecialForm("try", {REAL()});
  ASSERT_EQ(*tryResult, *REAL());
}

TEST_F(FunctionRegistryTest, resolveRowConstructor) {
  auto result = resolveFunctionOrCallableSpecialForm(
      "row_constructor", {INTEGER(), BOOLEAN(), DOUBLE()});
  ASSERT_EQ(
      *result, *ROW({"c1", "c2", "c3"}, {INTEGER(), BOOLEAN(), DOUBLE()}));
}

TEST_F(FunctionRegistryTest, resolveFunctionNotSpecialForm) {
  auto result = resolveFunctionOrCallableSpecialForm("func_one", {VARCHAR()});
  ASSERT_EQ(*result, *VARCHAR());
}

TEST_F(FunctionRegistryTest, resolveCast) {
  ASSERT_THROW(
      resolveFunctionOrCallableSpecialForm("cast", {VARCHAR()}),
      pollux::PolluxRuntimeError);
}

TEST_F(FunctionRegistryTest, resolveWithMetadata) {
  auto result = resolveFunctionWithMetadata("func_one", {VARCHAR()});
  EXPECT_TRUE(result.has_value());
  EXPECT_EQ(*result->first, *VARCHAR());
  EXPECT_TRUE(result->second.defaultNullBehavior);
  EXPECT_FALSE(result->second.deterministic);
  EXPECT_FALSE(result->second.supportsFlattening);

  result = resolveFunctionWithMetadata("func_two", {BIGINT(), INTEGER()});
  EXPECT_TRUE(result.has_value());
  EXPECT_EQ(*result->first, *BIGINT());
  EXPECT_FALSE(result->second.defaultNullBehavior);
  EXPECT_TRUE(result->second.deterministic);
  EXPECT_FALSE(result->second.supportsFlattening);

  result = resolveFunctionWithMetadata(
      "vector_func_four", {MAP(INTEGER(), VARCHAR())});
  EXPECT_TRUE(result.has_value());
  EXPECT_EQ(*result->first, *ARRAY(INTEGER()));
  EXPECT_TRUE(result->second.defaultNullBehavior);
  EXPECT_FALSE(result->second.deterministic);
  EXPECT_FALSE(result->second.supportsFlattening);

  result = resolveFunctionWithMetadata("non-existent-function", {VARCHAR()});
  EXPECT_FALSE(result.has_value());
}

class FunctionRegistryOverwriteTest : public functions::test::FunctionBaseTest {
 public:
  FunctionRegistryOverwriteTest() {
    registerTestFunctions();
  }
};

TEST_F(FunctionRegistryOverwriteTest, overwrite) {
  ASSERT_TRUE((register_function<FuncFive, int64_t, int64_t>({"foo"})));
  ASSERT_FALSE(
      (register_function<FuncSix, int64_t, int64_t>({"foo"}, {}, false)));
  ASSERT_TRUE((evaluateOnce<int64_t, int64_t>("foo(c0)", 0) == 5));
  ASSERT_TRUE((register_function<FuncSix, int64_t, int64_t>({"foo"})));
  ASSERT_TRUE((evaluateOnce<int64_t, int64_t>("foo(c0)", 0) == 6));

  auto& functions = exec::simple_functions();
  auto signatures = functions.get_function_signatures("foo");
  ASSERT_EQ(signatures.size(), 1);
}

TEST_F(FunctionRegistryTest, ipPrefixRegistration) {
  registerIPPrefixType();
  register_function<IPPrefixFunc, IPPrefix, IPPrefix>({"ipprefix_func"});

  auto& functions = exec::simple_functions();
  auto signatures = functions.get_function_signatures("ipprefix_func");
  ASSERT_EQ(signatures.size(), 1);

  auto result = resolveFunctionWithMetadata("ipprefix_func", {IPPREFIX()});
  EXPECT_TRUE(result.has_value());
  EXPECT_EQ(*result->first, *IPPREFIX());
  EXPECT_TRUE(result->second.defaultNullBehavior);
  EXPECT_TRUE(result->second.deterministic);
  EXPECT_FALSE(result->second.supportsFlattening);
}

} // namespace kumo::pollux
